package tuple.entity;

import org.yaml.snakeyaml.Yaml;
import soot.Body;
import soot.SootMethod;
import soot.Value;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

public class EntityRule {
    public static final Map<String, Pattern> ENTITY_PATTERNS = new HashMap<>();
    private static  String pathToEntityOntology =
//            "/home/fyl/project/code/AndroPrivacyConsistency/data/ontology/entity_ontology.yml";
        "C:\\Users\\Yalin Feng\\Code\\AndroidPrivacy\\Project\\data\\ontology\\entity_ontology.yml";

    private static boolean isLoaded = false;

    public static void ensureEntityLoaded() throws IOException {
        if (!isLoaded) {
            loadEntityFromYaml();
            isLoaded = true;
        }
    }

    public static void setEntityOntologyPath(String pathToEntityOntology){
        if (Files.exists(Paths.get(pathToEntityOntology))) {
            EntityRule.pathToEntityOntology = pathToEntityOntology;
        }
    }

    private static void loadEntityFromYaml() throws IOException {
        try (InputStream inputStream = Files.newInputStream(Paths.get(pathToEntityOntology))) {
            Yaml yaml = new Yaml();
            Map<String, Object> loadedYaml = yaml.load(inputStream);
            ENTITY_PATTERNS.putAll(parseLoaded(loadedYaml));
        }
    }

    private static Map<String, Pattern> parseLoaded(Map<String, Object> loadedYaml) {
        Map<String, Pattern> patterns = new HashMap<>();
        Queue<Map<String, Object>> queue = new ArrayDeque<>();
        queue.add(loadedYaml);
        while (!queue.isEmpty()) {
            Map<String, Object> curElement = queue.poll();
            if (curElement.containsKey("children")) {
                queue.addAll((Collection<? extends Map<String, Object>>) curElement.get("children"));
            }
            String name = (String) curElement.get("name");
            //patterns can String or ArrayList<String>
            if (curElement.get("patterns") instanceof ArrayList) {
                ArrayList list = (ArrayList) curElement.get("patterns");
                StringBuilder allPtn = null;
                for (Object o : list) {
                    if (o instanceof String s) {
                        if (allPtn == null)
                            allPtn = new StringBuilder(s);
                        else {
                            allPtn.append("|").append(s);
                        }
                    }
                }
                if (allPtn != null) {
                    patterns.put(name, Pattern.compile(allPtn.toString()));
                }
            } else if (curElement.get("patterns") instanceof String ptns) {
                Pattern p = Pattern.compile(ptns);
                patterns.put(name, p);
            }

        }
        return patterns;
    }

    public static Set<MyEntity> checkClassNameAsEntity(@Nonnull Set<String> splitClassName) {
        Set<MyEntity> entitySet = new HashSet<>();
        for (String classNamePart : splitClassName) {
            if (classNamePart.length() <= 1)
                continue;
            String classNamePartLowered = classNamePart.toLowerCase();
            for (String key : ENTITY_PATTERNS.keySet()) {
                Pattern pattern = ENTITY_PATTERNS.get(key);
                Matcher matcher = pattern.matcher(classNamePartLowered);
                if (matcher.find()) {
                    entitySet.add(MyEntity.valueOf(key));
                }
            }
        }
        return entitySet;
    }

    public static Set<MyEntity> checkMethodNameAsEntity(@Nonnull String name) {
        if (name.length() <= 1)
            return new HashSet<>();
        Set<MyEntity> entitySet = new HashSet<>();
        String nameLowered = name.toLowerCase();

        for (String key : ENTITY_PATTERNS.keySet()) {
            Pattern pattern = ENTITY_PATTERNS.get(key);
            Matcher matcher = pattern.matcher(nameLowered);
            if (matcher.find()) {
                entitySet.add(MyEntity.valueOf(key));
            }
        }
        return entitySet;
    }

    public static Set<MyEntity> checkMethodBodyAsEntity(@Nonnull Set<Value> noDupValues) {
        Set<MyEntity> entitySet = new HashSet<>();
        for (Value v : noDupValues) {
            String valLowered = v.toString().toLowerCase();
            for (String key : ENTITY_PATTERNS.keySet()) {
                Pattern pattern = ENTITY_PATTERNS.get(key);
                Matcher matcher = pattern.matcher(valLowered);
                if (matcher.find()) {
                    entitySet.add(MyEntity.valueOf(key));
                }
            }
        }
        return entitySet;
    }

    public static Set<MyEntity> checkMethodEntity(SootMethod sm) {
        Set<MyEntity> entitySet = new HashSet<>();

        // 1. className
        String className = sm.getDeclaringClass().getName();
        Set<String> splitName = Arrays.stream(className.split("\\."))
                .filter(x -> x.length() > 1)
                .collect(Collectors.toSet());
        entitySet.addAll(checkClassNameAsEntity(splitName));

        // 2. methodName
        entitySet.addAll(checkMethodNameAsEntity(sm.getSignature()));

        // 3. methodBody
        try {
            if (sm.hasActiveBody()) {
                Body body = sm.retrieveActiveBody();
                Set<Value> noDupValues = body.getUnits().stream()
                        .flatMap(x -> x.getUseAndDefBoxes().stream())
                        .map(x -> x.getValue())
                        .collect(Collectors.toSet());//JimpleLocal, JSpecialInvokeExpre,JNewExpr,JInstanceFieldRef,ThisRef
                entitySet.addAll(checkMethodBodyAsEntity(noDupValues));
            }
        } catch (RuntimeException re) {
            re.printStackTrace();
        }
        return entitySet;
    }
}
